from typing import Dict, List, Optional, Callable, Any
import pandas as pd
import numpy as np
from datetime import datetime
from itertools import product
import multiprocessing as mp
from dataclasses import dataclass
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns

@dataclass
class OptimizationResult:
    """优化结果"""
    params: Dict[str, Any]
    metrics: Dict[str, float]
    trades: pd.DataFrame
    
class StrategyOptimizer:
    """策略优化器"""
    
    def __init__(
        self,
        strategy_class: type,
        param_grid: Dict[str, List[Any]],
        market_data: pd.DataFrame,
        metric_functions: Dict[str, Callable] = None
    ):
        """
        初始化策略优化器
        
        Args:
            strategy_class: 策略类
            param_grid: 参数网格
            market_data: 市场数据
            metric_functions: 评估指标函数
        """
        self.strategy_class = strategy_class
        self.param_grid = param_grid
        self.market_data = market_data
        self.metric_functions = metric_functions or {
            'sharpe_ratio': self._calculate_sharpe_ratio,
            'max_drawdown': self._calculate_max_drawdown,
            'total_return': self._calculate_total_return,
            'win_rate': self._calculate_win_rate
        }
        self.results = []
        
    def _calculate_sharpe_ratio(self, trades: pd.DataFrame) -> float:
        """计算夏普比率"""
        if trades.empty:
            return 0.0
            
        daily_returns = trades['pnl'].resample('D').sum()
        if len(daily_returns) < 2:
            return 0.0
            
        return np.sqrt(252) * daily_returns.mean() / daily_returns.std()
        
    def _calculate_max_drawdown(self, trades: pd.DataFrame) -> float:
        """计算最大回撤"""
        if trades.empty:
            return 0.0
            
        cumulative_returns = (1 + trades['pnl']).cumprod()
        rolling_max = cumulative_returns.expanding().max()
        drawdowns = cumulative_returns / rolling_max - 1
        return drawdowns.min()
        
    def _calculate_total_return(self, trades: pd.DataFrame) -> float:
        """计算总收益"""
        if trades.empty:
            return 0.0
            
        return trades['pnl'].sum()
        
    def _calculate_win_rate(self, trades: pd.DataFrame) -> float:
        """计算胜率"""
        if trades.empty:
            return 0.0
            
        winning_trades = len(trades[trades['pnl'] > 0])
        return winning_trades / len(trades)
        
    def _run_backtest(self, params: Dict[str, Any]) -> OptimizationResult:
        """
        运行回测
        
        Args:
            params: 策略参数
            
        Returns:
            优化结果
        """
        # 创建策略实例
        strategy = self.strategy_class(**params)
        
        # 运行回测
        trades = []
        for i in range(len(self.market_data)):
            current_data = self.market_data.iloc[:i+1]
            strategy.update(current_data)
            
            # 收集交易记录
            if hasattr(strategy, 'trades'):
                trades.extend(strategy.trades)
                
        # 转换为DataFrame
        trades_df = pd.DataFrame([vars(t) for t in trades])
        if not trades_df.empty:
            trades_df['timestamp'] = pd.to_datetime(trades_df['timestamp'])
            trades_df.set_index('timestamp', inplace=True)
            
        # 计算评估指标
        metrics = {}
        for name, func in self.metric_functions.items():
            metrics[name] = func(trades_df)
            
        return OptimizationResult(params, metrics, trades_df)
        
    def optimize(self, n_jobs: int = -1) -> List[OptimizationResult]:
        """
        执行参数优化
        
        Args:
            n_jobs: 并行进程数
            
        Returns:
            优化结果列表
        """
        # 生成参数组合
        param_combinations = [
            dict(zip(self.param_grid.keys(), values))
            for values in product(*self.param_grid.values())
        ]
        
        # 并行执行回测
        with mp.Pool(n_jobs if n_jobs > 0 else mp.cpu_count()) as pool:
            self.results = list(tqdm(
                pool.imap(self._run_backtest, param_combinations),
                total=len(param_combinations)
            ))
            
        return self.results
        
    def plot_optimization_results(self, metric: str = 'sharpe_ratio', top_n: int = 10):
        """
        绘制优化结果
        
        Args:
            metric: 评估指标
            top_n: 显示前N个结果
        """
        if not self.results:
            return
            
        # 提取结果
        results_df = pd.DataFrame([
            {
                **r.params,
                metric: r.metrics[metric]
            }
            for r in self.results
        ])
        
        # 按指标排序
        results_df = results_df.sort_values(metric, ascending=False)
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle(f'Strategy Optimization Results - {metric}', fontsize=16)
        
        # 1. 参数分布
        for i, param in enumerate(self.param_grid.keys()):
            if i >= 4:
                break
            row, col = divmod(i, 2)
            sns.boxplot(
                x=param,
                y=metric,
                data=results_df,
                ax=axes[row, col]
            )
            axes[row, col].set_title(f'{param} vs {metric}')
            
        plt.tight_layout()
        plt.show()
        
        # 打印最佳结果
        print(f"\nTop {top_n} {metric} Results:")
        print(results_df.head(top_n))
        
    def get_best_params(self, metric: str = 'sharpe_ratio') -> Dict:
        """
        获取最佳参数
        
        Args:
            metric: 评估指标
            
        Returns:
            最佳参数
        """
        if not self.results:
            return {}
            
        # 按指标排序
        sorted_results = sorted(
            self.results,
            key=lambda x: x.metrics[metric],
            reverse=True
        )
        
        return sorted_results[0].params
